import numpy as np
from audio_processing.base_processor import AudioProcessor


class DelayProcessor(AudioProcessor):
    """延迟处理器"""

    def __init__(self, parameters):
        super().__init__(parameters)
        self.delay_time = parameters.get('delay_time', 0.5)  # 秒
        self.feedback = parameters.get('feedback', 0.3)
        self.mix = parameters.get('mix', 0.5)

        # 延迟缓冲区
        self.delay_buffer = None
        self.write_index = 0
        self.read_index = 0

    def set_sample_rate(self, sample_rate: int):
        super().set_sample_rate(sample_rate)
        self._init_delay_buffer()

    def set_input_channels(self, channels: int):
        """设置输入通道数并重新初始化缓冲区"""
        super().set_input_channels(channels)
        if self.sample_rate > 0:
            self._init_delay_buffer()

    def _init_delay_buffer(self):
        """初始化延迟缓冲区"""
        delay_samples = int(self.delay_time * self.sample_rate)
        if delay_samples > 0:
            # 根据输入通道数创建缓冲区
            self.delay_buffer = np.zeros((delay_samples, self.input_channels))
            self.write_index = 0
            self.read_index = (self.write_index - delay_samples) % len(self.delay_buffer)

    def process(self, audio_data: np.ndarray) -> np.ndarray:
        """应用延迟效果"""
        if not self.is_active or self.delay_buffer is None:
            return audio_data

        # 确保数据形状正确
        audio_data = self._ensure_correct_shape(audio_data)

        output = np.zeros_like(audio_data)

        for i in range(len(audio_data)):
            # 读取延迟样本
            delayed_sample = self.delay_buffer[self.read_index]

            # 混合原始音频和延迟音频
            output[i] = (1 - self.mix) * audio_data[i] + self.mix * delayed_sample

            # 写入当前样本加上反馈到延迟缓冲区
            self.delay_buffer[self.write_index] = audio_data[i] + self.feedback * delayed_sample

            # 更新索引
            self.write_index = (self.write_index + 1) % len(self.delay_buffer)
            self.read_index = (self.read_index + 1) % len(self.delay_buffer)

        return output

    def update_parameters(self, new_parameters):
        super().update_parameters(new_parameters)

        old_delay_time = self.delay_time
        self.delay_time = new_parameters.get('delay_time', self.delay_time)
        self.feedback = new_parameters.get('feedback', self.feedback)
        self.mix = new_parameters.get('mix', self.mix)

        # 如果延迟时间改变，重新初始化缓冲区
        if old_delay_time != self.delay_time and self.sample_rate > 0:
            self._init_delay_buffer()

    def reset(self):
        super().reset()
        if self.delay_buffer is not None:
            self.delay_buffer.fill(0)